from importlib.resources import path
import io
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image as PilImage
from paddle.io import Dataset
from paddle.vision.transforms import transforms as T
import random

IMAGE_SIZE = (720,720)

class CgDataSet(Dataset):
    """
    数据集定义
    """
    def __init__(self, filepath, mode='predict'):
        """
        构造函数
        """
        self.image_size = IMAGE_SIZE
        self.mode = mode.lower()
        self.path = filepath
        assert self.mode in ['predict'], \
            "mode should be 'predict', but got {}".format(self.mode)
        
        self.train_images = []
        # self.label_images = []

        self.train_images.append(filepath)
        
    def _load_img(self, color_mode='rgb', transforms=[]):
        """
        统一的图像处理接口封装，用于规整图像大小和通道
        """
        img = PilImage.open(self.path)
        if color_mode == 'grayscale':
            # if image is not already an 8-bit, 16-bit or 32-bit grayscale image
            # convert it to an 8-bit grayscale image.
            if img.mode not in ('L', 'I;16', 'I'):
                img = img.convert('L')
        elif color_mode == 'rgba':
            if img.mode != 'RGBA':
                img = img.convert('RGBA')
        elif color_mode == 'rgb':
            if img.mode != 'RGB':
                img = img.convert('RGB')
        else:
            raise ValueError('color_mode must be "grayscale", "rgb", or "rgba"')
        
        return T.Compose([
            T.Resize(self.image_size)
        ] + transforms)(img)



    def __getitem__(self, idx):
        """
        返回 image, label
        """
        train_image = self._load_img(   transforms=[
                                        T.Transpose(), 
                                        T.Normalize(mean=127.5, std=127.5)
                                     ]) # 加载原始图像
        # label_image = self._load_img(self.label_images[idx], 
        #                              color_mode='grayscale',
        #                              transforms=[T.Grayscale()]) # 加载Label图像
    
        # 返回image, label
        train_image = np.array(train_image, dtype='float32')
        # label_image = np.array(label_image, dtype='int64')
        return train_image
    def __len__(self):
        """
        返回数据集总数
        """
        return len(self.train_images)